"""
Created on Mar 1, 2020
Pytorch Implementation of LightGCN in
Xiangnan He et al. LightGCN: Simplifying and Powering Graph Convolution Network for Recommendation

@author: Shuxian Bi (stanbi@mail.ustc.edu.cn),Jianbai Ye (gusye@mail.ustc.edu.cn)
Design Dataset here
Every dataset's index has to start at 0
"""
import os
from os.path import join
import sys
import torch
import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from scipy.sparse import csr_matrix
import scipy.sparse as sp
import world
from world import cprint
from time import time

class BasicDataset(Dataset):
    def __init__(self):
        print("init dataset")
    #@property是一个装饰器 把一个方法变成属性 加上这个装饰器可以通过访问属性一样访问这个方法
    @property
    #获取所有用户数量 是一个数值
    def n_users(self):
        #如果子类没有重写这个方法，直接调用父类这个方法时会报错
        raise NotImplementedError
    
    @property
    #获取所有商品数量 是一个数值
    def m_items(self):
        raise NotImplementedError
    
    @property
    #获取训练集数据数量即所有用户的交互物品次数 是一个数值即统计交互矩阵中1的个数
    def trainDataSize(self):
        raise NotImplementedError
    
    @property
    #将交互矩阵以字典的形式表示 dict: {user: [items]} user为用户 [items]为该用户交互过的所有商品
    def testDict(self):
        raise NotImplementedError
    
    @property
    #返回用户交互过的商品 是一个列表 例allPos[i,:]获取第i行数据 即用户i所有交互过的商品列表
    def allPos(self):
        raise NotImplementedError
    #查看是否有交互
    def getUserItemFeedback(self, users, items):
        raise NotImplementedError
    #和allPos一样 但allPos是属性，getUserPosItems是方法
    def getUserPosItems(self, users):
        raise NotImplementedError
    #返回用户没有交互过的商品 是一个列表 例allPos[i,:]获取第i行数据 即用户i所有没有交互过的商品列表
    def getUserNegItems(self, users):
        """
        #不建议在大数据中返回用户没有交互过的商品 非常浪费性能
        not necessary for large dataset
        it's stupid to return all neg items in super large dataset
        """
        raise NotImplementedError
    """
    获得邻接矩阵并归一化，归一化即进行图卷积公式，会生成一个s_pre_adj_mat.npz来保存邻接矩阵
    最后读取邻接矩阵判断是否需要分割，如果不需要分割就把矩阵放入内存准备进行后续处理
    邻接矩阵维度A = 
            	|I,   R|
            	|R^T, I|
    """
    def getSparseGraph(self):
        """
        build a graph in torch.sparse.IntTensor.
        Details in NGCF's matrix form
        A = 
            |I,   R|
            |R^T, I|
        """
        raise NotImplementedError

class LastFM(BasicDataset):
    """
    Dataset type for pytorch \n
    Incldue graph information
    LastFM dataset
    """
    def __init__(self, path="../data/lastfm"):
        # train or test
        #自定义的cprint方法 用于改变字符颜色
        cprint("loading [last fm]")
        self.mode_dict = {'train':0, "test":1}
        #self.mode值为1
        self.mode    = self.mode_dict['train']
        # self.n_users = 1892
        # self.m_items = 4489
        #pd.read_table即pandas用于读取以制表符或其他符号分割的文件 header=None表明第一行作为数据而不是列名
        trainData = pd.read_table(join(path, 'data1.txt'), header=None)
        # print(trainData.head())
        testData  = pd.read_table(join(path, 'test1.txt'), header=None)
        # print(testData.head())
        trustNet  = pd.read_table(join(path, 'trustnetwork.txt'), header=None).to_numpy()
        # print(trustNet[:5])
        #把数据集中的所有元素都改为从0开始(一开始默认为1)
        trustNet -= 1
        trainData-= 1
        testData -= 1
        self.trustNet  = trustNet
        self.trainData = trainData
        self.testData  = testData
        #提取训练集所有用户的id 取训练数据每一行的第0列 np.array用于把输入数据转为numpy数组
        self.trainUser = np.array(trainData[:][0])
        #np.unique去掉重复元素只保留唯一一个
        self.trainUniqueUsers = np.unique(self.trainUser)
        #提取训练集所有商品ID
        self.trainItem = np.array(trainData[:][1])
        # self.trainDataSize = len(self.trainUser)
        #提取训练集所有用户的id
        self.testUser  = np.array(testData[:][0])
        #去重
        self.testUniqueUsers = np.unique(self.testUser)
        #提取训练集所有商品ID
        self.testItem  = np.array(testData[:][1])
        self.Graph = None
        print(f"LastFm Sparsity : {(len(self.trainUser) + len(self.testUser))/self.n_users/self.m_items}")
        
        # (users,users)
        """
        trustNet[:,0], trustNet[:,1]这部分是使用NumPy的切片操作来提取trustNet数组的第0列和第1列
        np.ones(len(trustNet))创建一个全是1的numpy类型 其维度只有一个即一个数组
        shape=(self.n_users,self.n_users) 形状为[self.n_users,self.n_users]
        这段话的意思是把trustNet[:,0]作为行，把trustNet[:,1])作为列，所有对应元素数据为1，结构为(self.n_users,self.n_users) 里面保存数据如下
        (0, 257)	1
        (0, 400)	1
        表示0号用户和257用户为朋友关系 0号用户和400用户为朋友关系
        """
        self.socialNet    = csr_matrix((np.ones(len(trustNet)), (trustNet[:,0], trustNet[:,1]) ), shape=(self.n_users,self.n_users))
        # (users,items), bipartite graph
        #建立用户和物品的稀疏矩阵
        self.UserItemNet  = csr_matrix((np.ones(len(self.trainUser)), (self.trainUser, self.trainItem) ), shape=(self.n_users,self.m_items)) 
        
        # range是从0到n_users-1的数组 getUserPosItems获得每一个用户交互过的物品数组 self._allPos获得每一个用户交互过的商品
        self._allPos = self.getUserPosItems(list(range(self.n_users)))
        #存储负样本
        self.allNeg = []
        #set集合不回出现重复的元素 range(self.m_items)返回一个0~range(self.m_items)-1的序列
        allItems    = set(range(self.m_items))
        #从0开始，到self.n_users - 1结束 [0,1892-1] 获取所有用户的负样本
        for i in range(self.n_users):
            #set给用户i交互过的所有商品去重
            pos = set(self._allPos[i])
            #把所用的商品集合减去用户交互过的商品集合相减得到负样本即用户没有交互过的集合
            neg = allItems - pos
            #将结果保存到self.allNeg数组中
            self.allNeg.append(np.array(list(neg)))
        # self.__build_test获得一个字典键为用户值为每个用户交互过的商品
        self.__testDict = self.__build_test()

    @property
    def n_users(self):
        return 1892
    
    @property
    def m_items(self):
        return 4489
    
    @property
    def trainDataSize(self):
        return len(self.trainUser)
    
    @property
    #返回一个字典键为用户值为每个用户交互过的商品
    def testDict(self):
        return self.__testDict

    @property
    #获得每一个用户交互过的商品
    def allPos(self):
        return self._allPos

    def getSparseGraph(self):
        if self.Graph is None:
            #self.trainUser[1,len_user]即[1,42135] 提取训练集所有用户的id没去重
            user_dim = torch.LongTensor(self.trainUser)
            #self.trainItem[1,len_item]即[1,42135] 提取训练集所有商品ID没去重
            item_dim = torch.LongTensor(self.trainItem)
            #item_dim + self.n_users用到python的广播 torch.stack将两个向量拼接 第二个向量接到第一个向量的下一行
            #first_sub[2,len_交互]即[2,42135]
            first_sub = torch.stack([user_dim, item_dim + self.n_users])
            #调换拼接次序 second_sub[2,len_交互]即[2,42135]
            second_sub = torch.stack([item_dim+ self.n_users, user_dim])
            #按列来拼接 即行数要求相等 index[2,84270]
            index = torch.cat([first_sub, second_sub], dim=1)
            #index倒数第一个维度即所有的用户和商品交互的总次数2*len_交互 创建一个对应长度为(2*len_交互)的一维数组元素全为1 [84270]
            data = torch.ones(index.size(-1)).int()
            """
            self.Graph维度[self.n_users+self.m_items, self.n_users+self.m_items]即[6381,6381]
            index [2,84270]保存了用户到物品的交互和物品到用户的交互 因为用户和物品的序号可能一样所以物品的序号是从用户序号的最大值开始的即在物品序号加上self.n_users
            torch.sparse.IntTensor：这是PyTorch中用于创建稀疏矩阵的函数
            data 84270 创建一个对应长度为(2*len_交互)的一维数组元素全为1用于表示用户和物品的交互
            torch.Size定义矩阵的大小
            """
            self.Graph = torch.sparse.IntTensor(index, data, torch.Size([self.n_users+self.m_items, self.n_users+self.m_items]))
            #dense的维度[6381,6381] Graph.to_dense()将稀疏矩阵转为正常矩阵
            dense = self.Graph.to_dense()
            #将dense矩阵按列相加[6381]
            D = torch.sum(dense, dim=1).float()
            #把矩阵D中所有等于0的元素替换为1 D即度矩阵 论文中为了防止分母为0将最小度默认为1
            D[D==0.] = 1.
            #D_sqrt维度为[1,6381]
            D_sqrt = torch.sqrt(D).unsqueeze(dim=0)
            #论文中所有元素和根号相除
            dense = dense/D_sqrt
            dense = dense/D_sqrt.t()
            #nonzero找到元素不为0的下标
            index = dense.nonzero()
            #选择dense中所有大于等于1e-9的元素 注意这里data的维度不是和dense维度一样，data是一个一维数组里面保存了所有>=1e-9的元素
            data  = dense[dense >= 1e-9]
            assert len(index) == len(data)
            """
            index.t() [2,84270] index:找到元素不为0的下标 一共84270对
            data 84270          data:
            graph [6381,6381]
            """
            self.Graph = torch.sparse.FloatTensor(index.t(), data, torch.Size([self.n_users+self.m_items, self.n_users+self.m_items]))
            #迁移到指定设备 coalesce()用于优化图存储结构
            self.Graph = self.Graph.coalesce().to(world.device)
        return self.Graph
    #获得一个字典键为用户值为每个用户交互过的商品
    def __build_test(self):
        """
        return:
            dict: {user: [items]}
        """
        test_data = {}
        #self.testItem 提取训练集所有商品ID没有去重
        for i, item in enumerate(self.testItem):
            # self.testUser 训练集所有用户的id没有去重
            user = self.testUser[i]
            #判断用户是否第一次加入字典
            if test_data.get(user):
                #不是第一次加入字典就把商品添加到末尾
                test_data[user].append(item)
            else:
                #是第一次加入就创建一个列表 第一个值就是item
                test_data[user] = [item]
        #最后获得一个字典键为用户值为每个用户交互过的商品
        return test_data
    
    def getUserItemFeedback(self, users, items):
        """
        users:
            shape [-1]
        items:
            shape [-1]
        return:
            feedback [-1]
        """
        # print(self.UserItemNet[users, items])
        return np.array(self.UserItemNet[users, items]).astype('uint8').reshape((-1, ))
    
    def getUserPosItems(self, users):
        posItems = []
        for user in users:
            #UserItemNet用户和物品的稀疏矩阵 nonzero()返回非0元素的维度信息 [1]取所有非0元素的第二个维度即下标为1的维度 append添加到数组末尾
            #用于获取一个用户所有的物品记录
            posItems.append(self.UserItemNet[user].nonzero()[1])
        return posItems
    
    def getUserNegItems(self, users):
        negItems = []
        for user in users:
            #self.allNeg存储负样本
            negItems.append(self.allNeg[user])
        return negItems
            
    
    #返回指定用户
    def __getitem__(self, index):
        #trainUniqueUsers去重后的用户
        user = self.trainUniqueUsers[index]
        # return user_id and the positive items of the user
        return user

    #切换到测试集
    def switch2test(self):
        """
        change dataset mode to offer test data to dataloader
        """
        self.mode = self.mode_dict['test']

    #返回去重后的用户长度
    def __len__(self):
        return len(self.trainUniqueUsers)

class Loader(BasicDataset):
    """
    Dataset type for pytorch \n
    Incldue graph information
    gowalla dataset
    """

    def __init__(self,config = world.config,path="../data/gowalla"):
        # train or test
        #打印路径
        cprint(f'loading [{path}]')
        #false
        self.split = config['A_split']
        #该参数用于指定如何分割大型邻接矩阵。如果没有提供该参数的值，其默认值为100
        self.folds = config['A_n_fold']
        self.mode_dict = {'train': 0, "test": 1}
        #选择训练集 即self.mode=0
        self.mode = self.mode_dict['train']
        self.n_user = 0
        self.m_item = 0
        train_file = path + '/train.txt'
        test_file = path + '/test.txt'
        self.path = path
        trainUniqueUsers, trainItem, trainUser = [], [], []
        testUniqueUsers, testItem, testUser = [], [], []
        self.traindataSize = 0
        self.testDataSize = 0
        #打开训练文件
        with open(train_file) as f:
            #遍历文件中的每一行。每次循环时，变量 l 包含文件中的一行
            for l in f.readlines():
                #当前行不为空时执行语句
                if len(l) > 0:
                    #使用 strip('\n') 方法移除行末的换行符，然后使用 split(' ') 方法将字符串按空格分割成一个列表
                    l = l.strip('\n').split(' ')
                    #从下标为1的元素即第二个元素开始每个元素转为int并赋值给items
                    items = [int(i) for i in l[1:]]
                    #把l第一个元素转为int给uid
                    uid = int(l[0])
                    #把uid添加到trainUniqueUsers数组中
                    trainUniqueUsers.append(uid)
                    #extend即给trainUser列表尾部加元素 元素内容为长度为len(items)的uid
                    #例 trainUser = [] ;  uid = 5 ;len(items)=3 执行下面代码后trainUser = [5,5,5]
                    trainUser.extend([uid] * len(items))
                    #把items列表添加到trainItem后
                    trainItem.extend(items)
                    """最后效果例 用户和商品一一对应
                        trainUser=[0,0,0,1,1,2,2,2,2....] 用户id
                        trainItem=[7,8,9,7,8,9,2,3,3....] 商品id
                    """
                    #找最大的商品编号
                    self.m_item = max(self.m_item, max(items))
                    #找最大的用户编号
                    self.n_user = max(self.n_user, uid)
                    #计算所有用户交互过的商品数量
                    self.traindataSize += len(items)
        #将下面三个数组转为numpy数组
        self.trainUniqueUsers = np.array(trainUniqueUsers)
        self.trainUser = np.array(trainUser)
        self.trainItem = np.array(trainItem)

        #和上面训练集操作一样
        with open(test_file) as f:
            # 遍历文件中的每一行。每次循环时，变量 l 包含文件中的一行
            for l in f.readlines():
                if len(l) > 0:
                    # 使用 strip('\n') 方法移除行末的换行符，然后使用 split(' ') 方法将字符串按空格分割成一个列表
                    l = l.strip('\n').split(' ')
                    # 从下标为1的元素即第二个元素开始每个元素转为int并赋值给items
                    items = [int(i) for i in l[1:]]
                    # 把l第一个元素转为int给uid
                    uid = int(l[0])
                    testUniqueUsers.append(uid)
                    testUser.extend([uid] * len(items))
                    testItem.extend(items)
                    self.m_item = max(self.m_item, max(items))
                    self.n_user = max(self.n_user, uid)
                    self.testDataSize += len(items)
        #m_item和m_item都是从0开始的 想要把他变成从1开始的 即1把最后结果加1
        self.n_user += 1
        self.m_item += 1
        self.testUniqueUsers = np.array(testUniqueUsers)
        self.testUser = np.array(testUser)
        self.testItem = np.array(testItem)
        
        self.Graph = None
        print(f"{self.trainDataSize} interactions for training")
        print(f"{self.testDataSize} interactions for testing")
        print(f"{world.dataset} Sparsity : {(self.trainDataSize + self.testDataSize) / self.n_users / self.m_items}")

        # (users,items), bipartite graph
        """
        np.ones(len(self.trainUser))创建一个长度为self.trainUser的全1向量 self.trainUser即富有用户交互商品的总次数
        self.trainUser 表示稀疏矩阵的非零元素的行索引。self.trainItem表示稀疏矩阵的非零元素的列索引。
        形状为(self.n_user, self.m_item)
        csr_matrix用于创建稀疏矩阵
        ../data/gowalla ：self.UserItemNet(self.n_user, self.m_item)即 (29858,40981)
        """
        self.UserItemNet = csr_matrix((np.ones(len(self.trainUser)), (self.trainUser, self.trainItem)),
                                      shape=(self.n_user, self.m_item))
        #对上面矩阵每一行相加axis=1即按列相加 例原式(3,4)变换后(3,1) 通过squeeze去掉维度为1的值
        #self.users_D计算每个用户和多少个物品进行交互
        self.users_D = np.array(self.UserItemNet.sum(axis=1)).squeeze()
        #将度为0的用户的度设置为1
        self.users_D[self.users_D == 0.] = 1
        #计算每个商品有多少用户交互
        self.items_D = np.array(self.UserItemNet.sum(axis=0)).squeeze()
        # 将度为0的商品的度设置为1
        self.items_D[self.items_D == 0.] = 1.
        # 获取用户交互过的所有商品 用户i的交互商品列表为 _allPos[i]
        self._allPos = self.getUserPosItems(list(range(self.n_user)))
        #__build_test()返回一个字典dict: {user: [items]} 用户和对应用户交互过的所有商品
        self.__testDict = self.__build_test()
        print(f"{world.dataset} is ready to go")

    @property
    def n_users(self):
        return self.n_user
    
    @property
    def m_items(self):
        return self.m_item
    
    @property
    def trainDataSize(self):
        return self.traindataSize
    
    @property
    def testDict(self):
        return self.__testDict

    @property
    def allPos(self):
        return self._allPos

    #它的目的是将一个矩阵 A 分割成多个子矩阵
    def _split_A_hat(self,A):
        A_fold = []
        #self.n_users 和 self.m_items 是用户和物品的数量 self.folds默认值为100 双除号//用于取结果的整数部分
        fold_len = (self.n_users + self.m_items) // self.folds
        #range(self.folds)即从[0,100)
        for i_fold in range(self.folds):
            #选择开始起点
            start = i_fold*fold_len
            #最后一次循环
            if i_fold == self.folds - 1:
                #选择结束 最后一次的结束是用户数加商品数
                end = self.n_users + self.m_items
            else:
                #选择结束
                end = (i_fold + 1) * fold_len
            #self._convert_sp_mat_to_sp_tensor稀疏矩阵格式转换转为torch类型
            A_fold.append(self._convert_sp_mat_to_sp_tensor(A[start:end]).coalesce().to(world.device))
        return A_fold
    #将一个来自 scipy.sparse 的稀疏矩阵转换为tensor的稀疏张量格式
    def _convert_sp_mat_to_sp_tensor(self, X):
        #tocoo() 方法将 X 转换为COO（Coordinate）格式，这是一种表示稀疏矩阵的格式。COO格式包含非零元素的行列索引和值
        coo = X.tocoo().astype(np.float32)
        #从COO格式中提取行索引，并将其转换为PyTorch的 long 类型张量。
        row = torch.Tensor(coo.row).long()
        #从COO格式中提取列索引，并将其转换为PyTorch的 long 类型张量。
        col = torch.Tensor(coo.col).long()
        #torch.stack按行堆叠 例 t1=([1,2,3])  t2=([4,5,6]) torch.stack(t1,t2)结果为tensor([[1, 2, 3],[14, 15, 16]])
        index = torch.stack([row, col])
        #从COO格式中提取非零元素的值，并将其转换为PyTorch的浮点数类型张量。
        data = torch.FloatTensor(coo.data)
        return torch.sparse.FloatTensor(index, data, torch.Size(coo.shape))
        
    def getSparseGraph(self):
        print("loading adjacency matrix")
        if self.Graph is None:
            try:
                #加载邻接矩阵
                pre_adj_mat = sp.load_npz(self.path + '/s_pre_adj_mat.npz')
                print("successfully loaded...")
                norm_adj = pre_adj_mat
            except :
                #没有预先加载的邻接矩阵
                print("generating adjacency matrix")
                s = time()
                #sp.dok_matrix用于创建一个“字典 of keys”类型的稀疏矩阵 当前稀疏矩阵中没有数据
                adj_mat = sp.dok_matrix((self.n_users + self.m_items, self.n_users + self.m_items), dtype=np.float32)
                #将字典形式转为 （List of Lists）格式，这是一个更紧凑的表示稀疏矩阵的方式，其中每个元素都是一个包含两个索引和一个值的元组。
                adj_mat = adj_mat.tolil()
                #把我们创建好的UserItemNet也转换为对应格式
                R = self.UserItemNet.tolil()
                #把格式转为论文中的格式
                adj_mat[:self.n_users, self.n_users:] = R
                adj_mat[self.n_users:, :self.n_users] = R.T
                #再把格式转回todok
                adj_mat = adj_mat.todok()
                # adj_mat = adj_mat + sp.eye(adj_mat.shape[0])
                #所有列相加结果为[row,1] rowsum[:self.n_users]为每个用户的度矩阵 self.n_users:为每个商品的度矩阵
                rowsum = np.array(adj_mat.sum(axis=1))
                #计算rowsum数组中每个元素的负0.5次方 flatten(): 这个方法将多维数组展平为一维数组。
                d_inv = np.power(rowsum, -0.5).flatten()
                #这个表达式将所有无穷大的值替换为0。这是为了防止在后续的计算中出现除以无穷大的错误。isinf判断是不是无穷大
                d_inv[np.isinf(d_inv)] = 0.
                #使用SciPy的diags函数来创建一个对角矩阵，其中对角线上的元素是d_inv中的值。
                d_mat = sp.diags(d_inv)
                #d_mat左乘adj_mat再右乘adj_mat
                #D^(-1/2) A D^(-1/2)
                norm_adj = d_mat.dot(adj_mat)
                norm_adj = norm_adj.dot(d_mat)
                #CSR是一种稀疏矩阵的存储格式，它特别适合于那些大部分元素都是零的矩阵
                norm_adj = norm_adj.tocsr()
                end = time()
                print(f"costing {end-s}s, saved norm_mat...")
                #将归一化后的邻接矩阵norm_adj保存为一个.npz文件。.npz文件是一种用于存储多个NumPy数组的文件格式。
                sp.save_npz(self.path + '/s_pre_adj_mat.npz', norm_adj)

            if self.split == True:
                #将图分割为多个小矩阵
                self.Graph = self._split_A_hat(norm_adj)
                print("done split matrix")
            else:
                #转换稀疏矩阵格式
                self.Graph = self._convert_sp_mat_to_sp_tensor(norm_adj)
                #运用gpu计算
                self.Graph = self.Graph.coalesce().to(world.device)
                print("don't split the matrix")
        return self.Graph
    #获取测试数据集 其格式为{user: [items]} 其中user为用户 [items]为用户交互过的所有商品序列
    def __build_test(self):
        """
        return:
            dict: {user: [items]}
        """
        test_data = {}
        #self.testItem所有用户交互过的商品列表    trainItem=[7,8,9,7,8,9,2,3,3....]
        for i, item in enumerate(self.testItem):
            #self.testUser交互过商品的所有用户列表trainUser=[0,0,0,1,1,2,2,2,2....]
            user = self.testUser[i]
            #在test_data字典中能不能找到编号为user的用户
            if test_data.get(user):
                #找到用户即在该字典下把item添加到数组尾部
                test_data[user].append(item)
            else:
                #找不到就新增一个数组
                test_data[user] = [item]
        return test_data

    def getUserItemFeedback(self, users, items):
        """
        users:
            shape [-1]
        items:
            shape [-1]
        return:
            feedback [-1]
        """
        # print(self.UserItemNet[users, items])
        return np.array(self.UserItemNet[users, items]).astype('uint8').reshape((-1,))

    #返回所有用户交互过的商品 返回 [users][所有商品交互信息]
    def getUserPosItems(self, users):
        posItems = []
        for user in users:
            """
            self.UserItemNet[user]返回用户物品交互图中对应用户的交互向量
            UserItemNet[user].nonzero()取交互向量中不为0的索引 返回两个列表第一个是用户索引 第二个是商品索引
            self.UserItemNet[user].nonzero()[1]获取所有的商品索引
            posItems.append 将该用户的索引添加到posItems列表的尾部
            """
            posItems.append(self.UserItemNet[user].nonzero()[1])
        return posItems

    # def getUserNegItems(self, users):
    #     negItems = []
    #     for user in users:
    #         negItems.append(self.allNeg[user])
    #     return negItems
